package com.zx.lc.rag;

import com.zx.lc.service.Assistant;
import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.DocumentSplitter;
import dev.langchain4j.data.document.loader.FileSystemDocumentLoader;
import dev.langchain4j.data.document.splitter.DocumentSplitters;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.ollama.OllamaChatModel;
import dev.langchain4j.model.ollama.OllamaEmbeddingModel;
import dev.langchain4j.model.output.Response;
import dev.langchain4j.rag.content.retriever.ContentRetriever;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.service.AiServices;
import dev.langchain4j.store.embedding.inmemory.InMemoryEmbeddingStore;
import java.time.Duration;
import java.util.List;
import org.springframework.stereotype.Service;

@Service
public class _2_NativeRagService {

    public static void main(String[] args) {
        //1. 初始化大语言模型，
        ChatLanguageModel chatModel = OllamaChatModel.builder()
                .baseUrl("http://192.168.3.99:11434")
                .timeout(Duration.ofSeconds(3600))
                .modelName("qwen2.5:7b")
                .build();

        // 2. 初始化嵌入模型：langchain4j支持的向量模型：https://docs.langchain4j.dev/category/embedding-models/
        OllamaEmbeddingModel embeddingModel = OllamaEmbeddingModel.builder()
                .baseUrl("http://192.168.3.99:11434")
                .modelName("bge-m3")
                .build();

        // 3.初始化向量数据库。使用内存模拟向量数据库
        // langchain4j支持的向量存储：https://docs.langchain4j.dev/integrations/embedding-stores/
        InMemoryEmbeddingStore<TextSegment> embeddingStore = new InMemoryEmbeddingStore<>();

        //  读取数据文件。
        //  1. 支持多种来源的数据解析：文件系统, URL, Amazon S3, Azure Blob存储, GitHub, 腾讯云 COS等
        //  2. 支持多种文件的解析：txt、pdf、doc、xls、ppt
        List<Document> documents = FileSystemDocumentLoader.loadDocumentsRecursively("D:\\llm\\rag-demo");

        // 将文件进行拆分为更小的块，方面精准的进行查询，而不是每次豆包珍整个文章传给模型，降低传输的数据量
        DocumentSplitter splitter = DocumentSplitters.recursive(300, 0);
        for (Document document : documents) {
            List<TextSegment> segments = splitter.split(document);
            // 使用嵌入模型，将拆分后的文件进行向量化
            Response<List<Embedding>> embeddings = embeddingModel.embedAll(segments);
            // 将向量化的数据存入向量数据库中
            embeddingStore.addAll(embeddings.content(), segments);
        }

        // 4. 初始化内容检索器，根据用户查询检索相关内容。
        // 目前，它能够检索文本段，但未来的增强功能将包括支持其他模态，如图像、音频等。
        ContentRetriever contentRetriever = EmbeddingStoreContentRetriever.builder()
                .embeddingStore(embeddingStore)
                .embeddingModel(embeddingModel)
                .maxResults(2) // 在每次交互中，我们将检索2个最相关的片段
                .minScore(0.5) // 向量计算相似度。最小相似度为0.5
                .build();

        Assistant assistant = AiServices.builder(Assistant.class)
                .chatLanguageModel(chatModel)
                .chatMemory(MessageWindowChatMemory.withMaxMessages(10))
                .contentRetriever(contentRetriever)
                .build();

        String chat = assistant.chat("陆瑶 是谁？帮我查下这个人的信息");
        System.out.println(chat);
    }


}
